"""
分析

1. 准备数据dataset
2. 初始化网络模型
3. 损失函数
   - 目标检测
     - 正样本
       - 置信度：二分类交叉熵
       - 坐标：均方差损失
       - 类别：交叉熵损失
     - 负样本
       - 置信度：二分类交叉熵

4. 优化器
"""
import os

import torch.optim
from torch import nn
from torch.utils.data import DataLoader
from yolov3 import YoLov3
from dataset import ODDateset
from config import cfg

device = cfg.device


class Train:
    def __init__(self):
        # 1. 准备数据dataset
        od_dataset = ODDateset()
        self.dataloader = DataLoader(od_dataset, batch_size=6, shuffle=True)
        # 2. 初始化网络模型
        self.net = YoLov3().to(device)
        # 加载参数
        # if os.path.exists(cfg.WEIGHT_PATH):
        #     self.net.load_state_dict(torch.load(cfg.WEIGHT_PATH))
        #     print('loading weights successfully')
        # 3. 损失函数
        # - 置信度：二分类交叉熵BCEWithLogitsLoss
        self.conf_loss_fn = nn.BCEWithLogitsLoss()
        # - 坐标：均方差损失MSELoss
        self.loc_loss_fn = nn.MSELoss()
        # - 类别：交叉熵损失CrossEntropyLoss
        self.cls_loss_fn = nn.CrossEntropyLoss()
        # 4. 优化器
        self.opt = torch.optim.Adam(self.net.parameters())

    def train(self, epoch):
        """
        :param epoch: 迭代训练的次数
        :return: None
        1. 开启训练
        2. 遍历数据加载器获取三种特征大小标签值、图片张量，并切换设备
        3. 图片张量传入网络获得三种预期输出
        4. 三种标签值、对应预期输出值和正负样本因子传入loss_fn，计算获得对应损失，并求和获得模型损失
        5. 优化器进行梯度清零
        6. 模型损失反向传播
        7. 优化器进行梯度更新
        8. 累加模型损失计算平均损失
        9. 保存模型权重
        """
        # 1. 开启训练
        self.net.train()
        # 累加损失
        sum_loss = 0
        for target_13, target_26, target_52, img in self.dataloader:
            # 2. 获取三种特征大小标签值、图片张量，并切换设备
            target_13, target_26, target_52 = target_13.to(device), target_26.to(device), target_52.to(device)
            img = img.to(device)
            # 3. 图片张量传入网络获得三种预期输出
            pred_out_13, pred_out_26, pred_out_52 = self.net(img)
            # 4.
            loss_13 = self.loss_fn(target_13, pred_out_13, scale_factor=cfg.SCALE_FACTOR_BIG)
            loss_26 = self.loss_fn(target_26, pred_out_26, scale_factor=cfg.SCALE_FACTOR_MID)
            loss_52 = self.loss_fn(target_52, pred_out_52, scale_factor=cfg.SCALE_FACTOR_SML)
            loss = loss_13 + loss_26 + loss_52
            # 5. 梯度清零
            self.opt.zero_grad()
            # 6. 反向传播
            loss.backward()
            # 7. 梯度更新
            self.opt.step()
            sum_loss += loss.item()

        avg_loss = sum_loss / len(self.dataloader)
        print(f'{epoch}\t{avg_loss}')
        if epoch % 10 == 0:
            print('save weight')
            torch.save(self.net.state_dict(), cfg.WEIGHT_PATH)

    def loss_fn(self, target, pre_out, scale_factor):
        """
        :param target: 标签
        :param pre_out: 预期输出
        :param scale_factor: 正负样本因子
        :return: 正负样本乘上对应规模因子的累加和
        1. 预期输出值更换通道 N 27 H W --> N H W 27 --> N H W 3 9
        2. 获取位置索引值
            - 正样本数据位置 target[..., 0] > 0
            - 负样本数据位置 target[..., 0] == 0
        3. 计算损失
            - 正样本：置信度 坐标 类别
            - 负样本：置信度
            - 索引获取
                - 0 置信度
                - 1:5 坐标
                - 5: 类别
        4. 正负样本乘上对应规模因子的累加和
        """
        # 1. 预期输出值更换通道 N 27 H W --> N H W 27 --> N H W 3 9
        pre_out = pre_out.permute((0, 2, 3, 1))
        n, h, w, _ = pre_out.shape
        pre_out = torch.reshape(pre_out, (n, h, w, 3, -1))
        # 2. 获取位置索引值 正样本数据位置 target[..., 0] > 0 负样本数据位置 target[..., 0] == 0
        mask_obj = target[..., 0] > 0
        mask_noobj = target[..., 0] == 0
        # 3. 计算损失
        # 正样本：置信度 坐标 类别
        target_obj = target[mask_obj]
        output_obj = pre_out[mask_obj]
        conf_loss = self.conf_loss_fn(output_obj[:, 0], target_obj[:, 0])
        loc_loss = self.loc_loss_fn(output_obj[:, 1:5], target_obj[:, 1:5])
        cls_loss = self.cls_loss_fn(output_obj[:, 5:], torch.argmax(target_obj[:, 5:], dim=1))
        loss_obj = conf_loss + loc_loss + cls_loss
        # 负样本：置信度
        target_noobj = target[mask_noobj]
        output_noobj = pre_out[mask_noobj]
        loss_noobj = self.conf_loss_fn(output_noobj[:, 0], target_noobj[:, 0])
        # 4. 正负样本乘上对应规模因子的累加和
        return loss_obj * scale_factor + loss_noobj * (1 - scale_factor)

    def run(self):
        for epoch in range(500):
            self.train(epoch)
        pass


if __name__ == '__main__':
    train = Train()
    train.run()
    pass
